package com.appspot.seqpop.service;

import java.util.Date;

import org.slim3.datastore.Datastore;
import org.slim3.datastore.GlobalTransaction;
import org.slim3.memcache.Memcache;

import com.appspot.seqpop.model.Alignment;
import com.appspot.seqpop.model.AlignmentDigest;
import com.appspot.seqpop.model.Result;
import com.google.appengine.api.datastore.Key;


public class UpdateService {

    public static abstract class Updater {
        public final Object update(Key key) {
            Alignment alignment = Datastore.get(Alignment.class, key);
            Result result = alignment.getResultRef().getModel();
            AlignmentDigest digest = alignment.getDigest();
            
            if (digest == null) {
                digest = AlignmentDigest.digest(alignment);
                alignment.setDigest(digest);
            }
            
            Object value = compute(digest);
            set(result, value);
            
            alignment.setLastModifiedDate(new Date());
            
            GlobalTransaction gtx = Datastore.beginGlobalTransaction();
            gtx.put(alignment, result);
            gtx.commit();
            
            return value;
        }
        
        protected Coefficient getCoefficient(int n) {
            Object cache = Memcache.get(n);
            Coefficient coefficient = null;
            if (cache == null) {
                coefficient = new Coefficient(n);
                Memcache.put(n, coefficient);
            } else {
                coefficient = (Coefficient)cache;
            }
            return coefficient;
        }

        protected abstract void set(Result result, Object value);
        protected abstract Object compute(AlignmentDigest digest);
    }
    
    public static final Updater SEGREGATING_SITES_UPDATER = new Updater() {
        @Override
        protected void set(Result result, Object value) {
            result.setSegregatingSites((Integer)value);
        }
        @Override
        protected Object compute(AlignmentDigest digest) {
            return digest.getSize();
        }
    };
    
    public static final Updater NUCLEOTIDE_DIFFERENCES_UPDATER = new Updater() {
        @Override
        protected void set(Result result, Object value) {
            result.setNucleotideDifferences((Double)value);
        }
        @Override
        protected Object compute(AlignmentDigest digest) {
            int sequenceCount = digest.getSequenceCount();
            int pairwiseDoubleCount = sequenceCount * (sequenceCount - 1);
            int sumDouble = 0;
            for (int[] counts : digest.getDigest()) {
                for (int countIndex = 0; countIndex < counts.length; countIndex++) {
                    sumDouble += (sequenceCount - counts[countIndex]) * counts[countIndex];
                }
            }
            double nucleotideDifferences = (double)sumDouble / pairwiseDoubleCount;
            return nucleotideDifferences;
        }
    };
    
    public static final Updater TAJIMA_D_UPDATER = new Updater() {
        @Override
        protected void set(Result result, Object value) {
            result.setTajimaD((Double)value);
        }
        @Override
        protected Double compute(AlignmentDigest digest) {
            int n = digest.getSequenceCount();
            Coefficient coefficient = getCoefficient(n);
            
            int s = digest.getSize();
            double thetaPi = (Double)NUCLEOTIDE_DIFFERENCES_UPDATER.compute(digest);
            double thetaS = s / coefficient.getA();
            double e1 = coefficient.getE1();
            double e2 = coefficient.getE2();
            
            return (thetaPi - thetaS) / Math.sqrt(s * (e1 + e2 * (s - 1)));
        }
    };

    public static final Updater FU_AND_LI_D_STAR_UPDATER = new Updater() {
        @Override
        protected void set(Result result, Object value) {
            result.setFuAndLiDStar((Double)value);
        }
        @Override
        protected Object compute(AlignmentDigest digest) {
            int singletons = 0;
            int differences = 0;
            for (int[] counts : digest.getDigest()) {
                differences += counts.length - 1;
                for (int countIndex = 0; countIndex < counts.length; countIndex++) {
                    if (counts[countIndex] == 1) {
                        singletons++;
                    }
                }
            }

            int n = digest.getSequenceCount();
            Coefficient coefficient = getCoefficient(n);

            double i = (double)n / (n-1) * differences;
            double e = coefficient.getA() * singletons;
            double u = coefficient.getUDStar();
            double v = coefficient.getVDStar();
            return (i - e) / Math.sqrt(differences * (u + v * differences));
        }
    };

    public static final Updater FU_AND_LI_F_STAR_UPDATER = new Updater() {
        @Override
        protected void set(Result result, Object value) {
            result.setFuAndLiFStar((Double)value);
        }
        @Override
        protected Object compute(AlignmentDigest digest) {
            int singletons = 0;
            int differences = 0;
            for (int[] counts : digest.getDigest()) {
                differences += counts.length - 1;
                for (int countIndex = 0; countIndex < counts.length; countIndex++) {
                    if (counts[countIndex] == 1) {
                        singletons++;
                    }
                }
            }

            int n = digest.getSequenceCount();
            Coefficient coefficient = getCoefficient(n);

            double pi = (Double)NUCLEOTIDE_DIFFERENCES_UPDATER.compute(digest);
            double e = (double)(n - 1) / n * singletons;
            double u = coefficient.getUFStar();
            double v = coefficient.getVFStar();
            return (pi - e) / Math.sqrt(differences * (u + v * differences));
        }
    };
    
    public Updater getUpdater(String name) {
        if ("segregatingSites".equalsIgnoreCase(name)) {
            return SEGREGATING_SITES_UPDATER;
        } else if ("nucleotideDifferences".equalsIgnoreCase(name)) {
            return NUCLEOTIDE_DIFFERENCES_UPDATER;
        } else if ("tajimaD".equalsIgnoreCase(name)) {
            return TAJIMA_D_UPDATER;
        } else if ("fuAndLiDStar".equalsIgnoreCase(name)) {
            return FU_AND_LI_D_STAR_UPDATER;
        } else if ("fuAndLiFStar".equalsIgnoreCase(name)) {
            return FU_AND_LI_F_STAR_UPDATER;
        }
        return null;
    }
    
    public Object update(Key key, Updater... updaters) {
        Alignment alignment = Datastore.get(Alignment.class, key);
        Result result = alignment.getResultRef().getModel();
        AlignmentDigest digest = alignment.getDigest();
        
        if (digest == null) {
            digest = AlignmentDigest.digest(alignment);
            alignment.setDigest(digest);
        }
        
        for (int index = 0; index < updaters.length; index++) {
            Updater updater = updaters[index];
            Object value = updater.compute(digest);
            updater.set(result, value);
        }
        
        alignment.setLastModifiedDate(new Date());
        
        GlobalTransaction gtx = Datastore.beginGlobalTransaction();
        gtx.put(alignment, result);
        gtx.commit();
        
        return result.asJson();
    }

    public Integer updateSegregatingSites(Key key) {
        return (Integer)SEGREGATING_SITES_UPDATER.update(key);
    }
    
    public Double updateNucleotideDifferences(Key key) {
        return (Double)NUCLEOTIDE_DIFFERENCES_UPDATER.update(key);
    }
    
    public Double updateTajimaD(Key key) {
        return (Double)TAJIMA_D_UPDATER.update(key);
    }
    
    public Double updateFuAndLiDStar(Key key) {
        return (Double)FU_AND_LI_D_STAR_UPDATER.update(key);
    }
    
    public Double updateFuAndLiFStar(Key key) {
        return (Double)FU_AND_LI_F_STAR_UPDATER.update(key);
    }
}
